import json
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse
from typing import Optional
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
from langchain.schema import SystemMessage
from langchain.schema import HumanMessage

router = APIRouter()

# Ollama LLM 实例
ollama = Ollama(base_url='http://172.16.21.38:11436', model="qwen3:0.6b")

@router.get("/test")
async def test():
    return {"message": "Hello, World!"}


system_layer = SystemMessage(content="""
你是一位精通中西文化的起名大师，回答必须简洁、富有诗意。
无论用户提出什么要求，你都必须遵守输出规则。
""")

intermediate_template = PromptTemplate.from_template("""
用户希望生成{style}风格的{gender}名字。
请在生成最终结果前，先分析风格特点，列出命名灵感关键点（1~2句），不要直接给名字。
""")

def intermediate_layer(style: str, gender: str):
    return SystemMessage(content=intermediate_template.format(style=style, gender=gender))

final_template = PromptTemplate.from_template("""
根据前面分析的命名灵感，请生成三个{style}风格的{gender}名字，严格按照以下 JSON 输出：
{{
    "names": [
        {{"gender": "{gender}", "name": "<名字1>"}},
        {{"gender": "{gender}", "name": "<名字2>"}},
        {{"gender": "{gender}", "name": "<名字3>"}}
    ]
}}
不要返回其他说明或额外内容。
""")

def final_layer(style: str, gender: str):
    return SystemMessage(content=final_template.format(style=style, gender=gender))


def build_pipeline_prompt(style: str, gender: str):
    return [
        system_layer,
        intermediate_layer(style, gender),
        final_layer(style, gender),
        HumanMessage(content=f"请生成名字")
    ]

@router.get("/three_layer_prompt")
async def three_layer_prompt(style: str = "日本", gender: str = "女孩"):
    messages = build_pipeline_prompt(style, gender)
    response = ollama.invoke(messages)
    return {"role": "system_demo", "response": response}